import random
import numpy as np
from network_topology import GBN
from const import TASK_NUM
'''
task:10*node_num,
pair:(start_point, end_point)

'''


class taskGenerator:
    def __init__(self, G):
        super().__init__()
        self.Graph = G
        self.node_num = G.nodes.__len__()

    # 返回了两种格式:
    # task_list:两组对应的list
    #    type:dict
    #    {‘tasks’:task_list, 'pairs':pair_list}
    #
    # task_dict:可以根据source_target随机访问
    #    type:dict
    #    key:(start, end)
    #    value abs(task_payload)

    def generate(self, floor, ceil):
        task_matrix = np.zeros((TASK_NUM, self.node_num))
        pairs = []
        payloads = {}
        for i in range(0, TASK_NUM):
            start = random.randint(0, self.node_num - 1)
            end = random.randint(0, self.node_num - 1)
            while start == end:
                end = random.randint(0, self.node_num - 1)
            package = random.randint(floor, ceil)
            '''
            test
            '''
            # task_matrix[i, :] = random.randint(50, 100)
            task_matrix[i, start] = -package
            task_matrix[i, end] = package
            pairs.append((start, end))
            payloads[(start, end)] = package
        return {'task_matrix': task_matrix, 'pairs': pairs}, payloads

    def generate_group(self, group_num, floor, ceil):
        group = []
        for i in range(0, group_num):
            task_lists, task_dic = self.generate(floor, ceil)
            group.append({'task_lists': task_lists, 'task_dic': task_dic})
        return group


'''
G = GBN()
generator = taskGenerator(G)
print(generator.generate(40, 70))

batch = generator.generate_batch(5, 10, 20)
print(batch[4])
'''
